"""
Demo 05: Shapelet 距离特征
使用判别性子序列（shapelet）提取时间序列特征
注意：这个演示使用简化的shapelet实现，实际的ShapeletTransform计算量较大
"""

import numpy as np
from demo_01_raw import RawVectorizer
try:
    from tslearn.shapelets import ShapeletModel
except ImportError:
    print("请先安装tslearn: pip install tslearn")
    ShapeletModel = None


class ShapeletVectorizer:
    """Shapelet距离特征向量化处理类"""
    
    def __init__(self, window_size=60, step_size=5, n_shapelets=5, shapelet_length=30):
        """
        初始化参数
        
        Args:
            window_size: 窗口大小
            step_size: 步长
            n_shapelets: shapelet数量
            shapelet_length: 每个shapelet的长度
        """
        self.window_size = window_size
        self.step_size = step_size
        self.n_shapelets = n_shapelets
        self.shapelet_length = shapelet_length
        self.raw_vectorizer = RawVectorizer(window_size, step_size)
        self.shapelets = None
        
    def extract_random_shapelets(self, windows):
        """
        随机提取shapelet（简化版本）
        
        Args:
            windows: 窗口矩阵
            
        Returns:
            numpy.ndarray: shapelet矩阵
        """
        n_windows = windows.shape[0]
        shapelets = []
        
        for _ in range(self.n_shapelets):
            # 随机选择一个窗口
            window_idx = np.random.randint(0, n_windows)
            # 随机选择起始位置
            max_start = self.window_size - self.shapelet_length
            start_idx = np.random.randint(0, max_start + 1)
            # 提取子序列作为shapelet
            shapelet = windows[window_idx, start_idx:start_idx + self.shapelet_length]
            shapelets.append(shapelet)
        
        return np.array(shapelets)
    
    def compute_shapelet_distance(self, series, shapelet):
        """
        计算时间序列与shapelet的最小距离
        
        Args:
            series: 时间序列
            shapelet: shapelet子序列
            
        Returns:
            float: 最小距离
        """
        min_dist = np.inf
        shapelet_len = len(shapelet)
        
        # 滑动窗口计算距离
        for i in range(len(series) - shapelet_len + 1):
            subseries = series[i:i + shapelet_len]
            dist = np.sqrt(np.sum((subseries - shapelet) ** 2))
            min_dist = min(min_dist, dist)
        
        return min_dist
    
    def transform_with_shapelets(self, windows, shapelets):
        """
        使用shapelets转换窗口为距离特征
        
        Args:
            windows: 窗口矩阵
            shapelets: shapelet矩阵
            
        Returns:
            numpy.ndarray: 距离特征矩阵
        """
        n_windows = windows.shape[0]
        n_shapelets = shapelets.shape[0]
        distances = np.zeros((n_windows, n_shapelets))
        
        for i in range(n_windows):
            for j in range(n_shapelets):
                distances[i, j] = self.compute_shapelet_distance(
                    windows[i], shapelets[j]
                )
        
        return distances
    
    def extract_shapelet_statistics(self, distances):
        """
        从shapelet距离中提取统计特征
        
        Args:
            distances: 距离矩阵
            
        Returns:
            numpy.ndarray: 统计特征矩阵
        """
        # 计算距离的统计量
        mean_dist = distances.mean(axis=1, keepdims=True)
        std_dist = distances.std(axis=1, keepdims=True)
        min_dist = distances.min(axis=1, keepdims=True)
        max_dist = distances.max(axis=1, keepdims=True)
        
        # 找出最近的shapelet索引（one-hot编码）
        nearest_shapelet = np.zeros((distances.shape[0], self.n_shapelets))
        nearest_idx = distances.argmin(axis=1)
        for i, idx in enumerate(nearest_idx):
            nearest_shapelet[i, idx] = 1
        
        # 合并特征
        features = np.hstack([
            distances,  # 原始距离
            mean_dist, std_dist, min_dist, max_dist,  # 统计量
            nearest_shapelet  # 最近shapelet指示
        ])
        
        return features
    
    def fit_transform(self, ticker="AAPL", start="2020-01-01", end="2024-12-31"):
        """
        完整的Shapelet特征向量化流程
        
        Args:
            ticker: 股票代码
            start: 开始日期
            end: 结束日期
            
        Returns:
            tuple: (原始窗口矩阵, shapelet距离矩阵, 扩展特征矩阵)
        """
        # 获取原始窗口
        X_raw = self.raw_vectorizer.fit_transform(ticker, start, end)
        
        # 提取shapelets
        self.shapelets = self.extract_random_shapelets(X_raw)
        
        # 计算距离特征
        X_shapelet = self.transform_with_shapelets(X_raw, self.shapelets)
        
        # 提取扩展特征
        X_shapelet_ext = self.extract_shapelet_statistics(X_shapelet)
        
        return X_raw, X_shapelet, X_shapelet_ext
    
    def find_best_matches(self, distances, top_k=5):
        """
        找出与每个shapelet最匹配的窗口
        
        Args:
            distances: 距离矩阵
            top_k: 返回前k个匹配
            
        Returns:
            dict: 每个shapelet的最佳匹配
        """
        best_matches = {}
        
        for j in range(self.n_shapelets):
            shapelet_distances = distances[:, j]
            top_indices = np.argsort(shapelet_distances)[:top_k]
            best_matches[f"shapelet_{j}"] = {
                "indices": top_indices.tolist(),
                "distances": shapelet_distances[top_indices].tolist()
            }
        
        return best_matches


def main():
    """主函数，演示使用方法"""
    print("=" * 60)
    print("Demo 05: Shapelet 距离特征")
    print("=" * 60)
    
    # 创建Shapelet向量化器
    vectorizer = ShapeletVectorizer(
        window_size=60, 
        step_size=5, 
        n_shapelets=5, 
        shapelet_length=30
    )
    
    # 执行向量化
    print("正在下载AAPL股票数据并提取Shapelet特征...")
    print("（注意：Shapelet计算可能需要一些时间）")
    X_raw, X_shapelet, X_shapelet_ext = vectorizer.fit_transform(
        ticker="AAPL",
        start="2020-01-01",
        end="2024-12-31"
    )
    
    # 输出结果
    print(f"\n原始窗口 shape: {X_raw.shape}")
    print(f"Shapelet距离向量 shape: {X_shapelet.shape}")
    print(f"扩展特征向量 shape: {X_shapelet_ext.shape}")
    
    # 显示第一个窗口的距离特征
    print(f"\n第一个窗口到各个shapelet的距离:")
    for i in range(vectorizer.n_shapelets):
        print(f"  Shapelet {i}: {X_shapelet[0, i]:.4f}")
    
    # 统计距离分布
    print(f"\n距离统计:")
    print(f"  平均距离: {X_shapelet.mean():.4f}")
    print(f"  最小距离: {X_shapelet.min():.4f}")
    print(f"  最大距离: {X_shapelet.max():.4f}")
    
    # 找出每个shapelet的最佳匹配
    best_matches = vectorizer.find_best_matches(X_shapelet, top_k=3)
    print(f"\n每个Shapelet的最佳匹配窗口（前3个）:")
    for shapelet_name, matches in best_matches.items():
        indices = matches["indices"]
        distances = matches["distances"]
        print(f"  {shapelet_name}: 窗口 {indices} (距离: {[f'{d:.3f}' for d in distances]})")
    
    # 分析shapelet的判别性
    print(f"\nShapelet判别性分析:")
    for i in range(vectorizer.n_shapelets):
        distances_i = X_shapelet[:, i]
        print(f"  Shapelet {i}: 均值={distances_i.mean():.3f}, "
              f"标准差={distances_i.std():.3f}, "
              f"变异系数={distances_i.std()/distances_i.mean():.3f}")
    
    return X_raw, X_shapelet, X_shapelet_ext


if __name__ == "__main__":
    X_raw, X_shapelet, X_shapelet_ext = main()
